import socket
import threading
import sys

class PyClient:
    def __init__(self, host, port):
        self.host = host
        self.port = port
        self.client_socket = None
        self.running = False
        self.lock = threading.Lock()

    def connect(self):
        try:
            # 创建TCP socket
            self.client_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            # 连接服务器
            self.client_socket.connect((self.host, self.port))
            self.running = True
            print(f"成功连接到服务器 {self.host}:{self.port}")
            return True
        except Exception as e:
            print(f"连接失败: {str(e)}")
            return False

    def receive_handler(self):
        buffer = b''  # 新增：二进制数据缓冲区
        while self.running:
            try:
                data = self.client_socket.recv(4096)
                if not data:
                    self.stop()
                    print("服务器断开连接")
                    break
                
                buffer += data  # 累积数据
                # 尝试解码已有数据
                while len(buffer) > 0:
                    try:
                        # 将缓冲区数据转为字符串
                        message = buffer.decode('utf-8')
                        print(f"\n[来自服务器]: {message}")
                        buffer = b''  # 清空缓冲区
                        break
                    except UnicodeDecodeError:
                        # 如果解码失败，说明数据不完整，继续接收
                        break
                    
                print("> ", end="", flush=True)
            except Exception as e:
                if self.running:
                    print(f"\n接收错误: {str(e)}")
                self.stop()

    def send_handler(self):
        try:
            while self.running:
                # 获取用户输入
                message = input("> ")
                if not self.running:
                    break
                if message.lower() == "exit":
                    self.stop()
                    break
                try:
                    # 发送消息到服务器
                    self.client_socket.sendall(message.encode("utf-8"))
                except Exception as e:
                    print(f"发送失败: {str(e)}")
                    self.stop()
        except KeyboardInterrupt:
            self.stop()

    def stop(self):
        with self.lock:
            if self.running:
                self.running = False
                try:
                    if self.client_socket:
                        self.client_socket.close()
                except Exception as e:
                    pass
                print("连接已关闭")

    def start(self):
        if not self.connect():
            return

        # 启动接收线程
        receive_thread = threading.Thread(target=self.receive_handler, daemon=True)
        receive_thread.start()

        # 在主线程处理发送
        self.send_handler()

        # 等待接收线程结束
        receive_thread.join(2)

def main():
    if len(sys.argv) != 3:
        print("使用方法: python client.py <服务器IP> <端口>")
        print("示例: python client.py 123.45.67.89 8080")
        return

    host = sys.argv[1]
    port = int(sys.argv[2])

    client = PyClient(host, port)
    client.start()

if __name__ == "__main__":
    main()